//! low-level code to handle traps from user space into the kernel, and returns from kernel to user.
//!
//! the kernel maps the page holding this code at the same virtual address (TRAMPOLINE) in user and
//! kernel space so that it continues to work when it switches page tables. kernel.ld causes this code
//! to start at a page boundary.

const std = @import("std");
const print = std.fmt.comptimePrint;
const os = @import("../os.zig");

/// defined by kernel.ld
pub extern const _trampoline: u8;
const TRAPFRAME = os.mem.VirtAddr.TRAPFRAME;
const TRAMPOLINE = os.mem.VirtAddr.TRAMPOLINE;

/// per-process data for the trap handling code.
/// sits in a page by itself just under the trampoline page in the user page table. not specially mapped
/// in the kernel page table. uservec saves user registers in the trapframe, then initializes registers
/// from the trapframe's sp, hartid, satp, and jumps to trap. usertrapret() and userret set up the
/// trapframe, restore user registers from the trapframe, switch to the user page table, and enter user
/// space.
pub const TrapFrame = extern struct {
    /// kernel page table
    satp: usize,
    /// top of process's kernel stack
    sp: usize,
    /// usertrap()
    trap: *const fn () callconv(.c) noreturn,
    /// saved user program counter
    epc: usize,
    /// saved kernel tp
    hartid: usize,

    // callee-saved user registers
    reg: Registers,

    /// the trapframe includes callee-saved user registers like s0-s11 because the return-to-user path
    /// via usertrapret() doesn't return through the entire kernel call stack.
    pub const Registers = extern struct {
        ra: usize,
        sp: usize,
        gp: usize,
        tp: usize,
        t0: usize,
        t1: usize,
        t2: usize,
        s0: usize,
        s1: usize,
        a0: usize,
        a1: usize,
        a2: usize,
        a3: usize,
        a4: usize,
        a5: usize,
        a6: usize,
        a7: usize,
        s2: usize,
        s3: usize,
        s4: usize,
        s5: usize,
        s6: usize,
        s7: usize,
        s8: usize,
        s9: usize,
        s10: usize,
        s11: usize,
        t3: usize,
        t4: usize,
        t5: usize,
        t6: usize,
    };

    /// Enum type for register field access
    const Register = std.meta.FieldEnum(Registers);

    /// List of all register names in Registers
    const registers = std.meta.tags(Register);

    /// Returns the byte offset of a register within TrapFrame
    pub inline fn regOffset(comptime reg: Register) usize {
        return @offsetOf(TrapFrame, "reg") + @offsetOf(Registers, @tagName(reg));
    }
};

/// trap.zig sets stvec to point here, so traps from user space start here, in supervisor mode, but with
/// a user page table.
pub fn uservec() align(4) linksection("trampsec") callconv(.naked) noreturn {
    // save the first register in sscratch so it can be used to get at TRAPFRAME.
    const first = TrapFrame.registers[0];
    asm volatile (print("csrw sscratch, {t}", .{first}));

    // each process has a separate trapframe memory area, but it's mapped to the same virtual address
    // (TRAPFRAME) in every process's user page table.

    // save the user registers except the first register in TRAPFRAME
    asm volatile (print("li {t}, {d}", .{ first, TRAPFRAME }));
    inline for (TrapFrame.registers[1..]) |register| {
        os.arch.op(.save, register, TrapFrame.regOffset(register), first);
    }

    // save the first register
    const tmp = .t0;
    comptime std.debug.assert(first != tmp);
    asm volatile (print("csrr {t}, sscratch", .{tmp}));
    os.arch.op(.save, tmp, TrapFrame.regOffset(first), first);

    // loads different fields from the TrapFrame structure into registers
    inline for (.{ "sp", "hartid", "trap", "satp" }, .{ .sp, .tp, .t0, .t1 }) |field, reg| {
        os.arch.op(.load, reg, @offsetOf(TrapFrame, field), first);
    }

    // wait for any previous memory operations to complete, so that they use the user page table.
    asm volatile ("sfence.vma zero, zero");
    // install the kernel page table.
    asm volatile ("csrw satp, t1");
    // flush now-stale user entries from the TLB.
    asm volatile ("sfence.vma zero, zero");

    // jump to usertrap(), which does not return
    asm volatile ("jr t0");
}

test "trap-vector alignment" {
    const uservec_addr = @intFromPtr(&uservec);
    const trampoline_addr = @intFromPtr(&_trampoline);
    const trampoline_uservec = TRAMPOLINE + (uservec_addr - trampoline_addr);
    try std.testing.expect(std.mem.isAligned(trampoline_uservec, 4));
}

/// fn userret(satp: usize) callconv(.c) noreturn;
/// called by usertrapret() in trap.zig to switch from kernel to user.
pub fn userret() linksection("trampsec") callconv(.naked) noreturn {
    // switch to the user page table.
    asm volatile (
        \\ sfence.vma zero, zero
        \\ csrw satp, a0
        \\ sfence.vma zero, zero
    );

    // restore all from TRAPFRAME
    const last = TrapFrame.registers[TrapFrame.registers.len - 1];
    asm volatile (print("li {t}, {d}", .{ last, TRAPFRAME }));

    inline for (TrapFrame.registers) |reg| {
        os.arch.op(.load, reg, TrapFrame.regOffset(reg), last);
    }

    // return to user mode and user pc. usertrapret() set up sstatus and sepc.
    asm volatile ("sret");
}
